"""pregel.py is a python module implementing a toy single-machine
version of Google's Pregel system for large-scale graph processing."""

import collections
import threading
from enum import Enum

from abc import abstractmethod


class Vertex:

    def __init__(self, vid, value, out_vertices: list = []):
        # This is mostly self-explanatory, but has a few quirks:
        #
        # self.id is included mainly because it's described in the
        # Pregel paper.  It is used briefly in the pagerank example,
        # but not in any essential way, and I was tempted to omit it.
        #
        # Each vertex stores the current superstep number in
        # self.superstep.  It's arguably not wise to store many copies
        # of global state in instance variables, but Pregel's
        # synchronous nature lets us get away with it.
        self.id = vid
        self.value = value
        self.out_vertices = out_vertices
        self.incoming_messages = []
        self.outgoing_messages = []
        self.active = True
        self.superstep = 0

    def getid(self):
        return self.id

    def setValue(self, value):
        self.value = value

    def getValue(self):
        return self.value

    def send(self, target_v, message):
        self.outgoing_messages.append((target_v, message))

    def send2All(self, message):
        self.outgoing_messages = [(target_v, message) for target_v in self.out_vertices]

    @abstractmethod
    def compute(self, messages, context: dict = None):
        pass

    @abstractmethod
    def init(self, context: dict):
        # Initialize and send first message here, or just halt itself
        pass

    def set_out_vertices(self, out_vertices: list):
        self.out_vertices=out_vertices

    def get_out_vertices(self) -> list:
        return self.out_vertices

    def vote_to_halt(self):
        self.active = False

    def is_active(self) -> bool:
        return self.active


class Pregel:

    def __init__(self, num_workers):
        self.vertices = []
        self.vertices_builder={}
        self.num_workers = num_workers

    def set_Vertex(self, vbuilder):
        self.vertices_builder[vbuilder.vid] = vbuilder
        for adjid in vbuilder.out_vertices_id:
            if adjid not in self.vertices_builder:
                self.vertices_builder[adjid]=None

    def construct(self):
        '''
        检查每个出边是否都有对应实体，然后构建相应顶点
        '''
        if None in self.vertices_builder.values():
            raise RuntimeError("Graph has incomplete vertex")
        else:
            v_dict=dict()
            idx=0
            for vid,builder in self.vertices_builder.items():
                self.vertices.append(builder.build())
                v_dict[vid] = idx
                idx+=1
            for vid,builder in self.vertices_builder.items():
                out_vertices=[]
                for adjid in builder.out_vertices_id:
                    out_vertices.append(self.vertices[v_dict[adjid]])
                self.vertices[v_dict[vid]].set_out_vertices(out_vertices)

    def run(self, context=None,limit_step=50):
        """Runs the Pregel instance."""
        self.construct()
        self.partition = self.partition_vertices()
        self.initstep(context)
        now_step=0
        while now_step<limit_step and self.check_active():
            self.superstep(context)
            self.redistribute_messages()
            now_step+=1

    def partition_vertices(self):
        """Returns a dict with keys 0,...,self.num_workers-1
        representing the worker threads.  The corresponding values are
        lists of vertices assigned to that worker."""
        partition = collections.defaultdict(list)
        for vertex in self.vertices:
            partition[self.worker(vertex)].append(vertex)
        return partition

    def worker(self, vertex):
        """Returns the id of the worker that vertex is assigned to."""
        return hash(vertex) % self.num_workers

    def superstep(self, context):
        """Completes a single superstep.

        Note that in this implementation, worker threads are spawned,
        and then destroyed during each superstep.  This creation and
        destruction causes some overhead, and it would be better to
        make the workers persistent, and to use a locking mechanism to
        synchronize.  The Pregel paper suggests that this is how
        Google's Pregel implementation works.
        线程暂停: threading.Event
        """
        workers = []
        for vertex_list in self.partition.values():
            worker = Worker(vertex_list, context,Worker.commission_code.SUPERSTEP)
            workers.append(worker)
            worker.start()
        for worker in workers:
            worker.join()

    def initstep(self,context):
        workers = []
        for vertex_list in self.partition.values():
            worker = Worker(vertex_list, context, Worker.commission_code.INIT)
            workers.append(worker)
            worker.start()
        for worker in workers:
            worker.join()

    def redistribute_messages(self):
        """Updates the message lists for all vertices."""
        for vertex in self.vertices:
            vertex.superstep += 1
            vertex.incoming_messages = []
        for vertex in self.vertices:
            for receiving_vertix, message in vertex.outgoing_messages:
                receiving_vertix.incoming_messages.append((vertex, message))
                receiving_vertix.active=True
            vertex.outgoing_messages=[]

    def check_active(self):
        """Returns True if there are any active vertices, and False
        otherwise."""
        return any([vertex.active for vertex in self.vertices])


class Worker(threading.Thread):
    class commission_code(Enum):
        INIT=1
        SUPERSTEP=2

    def __init__(self, vertices: Vertex, context: dict, code: commission_code):
        threading.Thread.__init__(self)
        self.vertices = vertices
        self.context = context
        self.code = code

    def run(self):
        if self.code ==self.commission_code.SUPERSTEP:
            self.superstep()
        elif self.code ==self.commission_code.INIT:
            self.initstep()

    def initstep(self):
        for vertex in self.vertices:
            vertex.init(self.context)

    def superstep(self):
        """Completes a single superstep for all the vertices in
        self."""
        for vertex in self.vertices:
            if vertex.is_active():
                vertex.compute(vertex.incoming_messages, self.context)
